import time

from ultralytics import YOLO

import cv2
import numpy as np

class CLAHEPreprocessor:
    def __init__(self, clip_limit=2.0, tile_grid_size=(8, 8), brightness_balance=None):
        """
        初始化 CLAHE 预处理器
        :param clip_limit: CLAHE算法中的限制对比度因子，默认为2.0
        :param tile_grid_size: CLAHE算法中的网格大小，默认为8x8
        :param brightness_balance: 亮度平衡因子，如果为None，则不进行亮度平衡
        """
        self.clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
        self.brightness_balance = brightness_balance

    def __call__(self, left_image, right_image):
        """
        读取左右眼图像，并应用 CLAHE 进行对比度增强、亮度平衡和去噪
        :param left_image: 左眼图像
        :param right_image: 右眼图像
        :return: 返回增强后的左右眼图像
        """
        # 对左眼图像进行处理
        left_img = self._process_image(left_image)
        # 对右眼图像进行处理
        right_img = self._process_image(right_image)

        return left_img, right_img

    def _process_image(self, image):
        """
        对单张图像进行处理，包括CLAHE对比度增强、亮度平衡和去噪
        :param image: 输入图像
        :return: 处理后的图像
        """
        # 如果设置了亮度平衡因子，则调整图像亮度
        if self.brightness_balance:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            brightness = gray.sum() / gray.size
            image = np.uint8(np.minimum(image * self.brightness_balance / brightness, 255))

        # 应用 CLAHE 对比度增强
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        lab_planes = list(cv2.split(lab))  # 将元组转换为列表
        lab_planes[0] = self.clahe.apply(lab_planes[0])  # 对亮度通道进行CLAHE处理
        lab = cv2.merge(lab_planes)  # 重新组合为图像
        image = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

        # 应用去噪
        image = cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 1, 3)
        image = cv2.bilateralFilter(image, 5, 1, 1)

        return image
class EyeDiseasePredictor:
    """
    眼底疾病预测器
    输入双目图像路径/直接输入图像（np数组格式），输出一个字典
    字典包含左眼和右眼的预测结果，包含以下信息：
    - Left Eye:
        -- result:
            -- Predicted Class ID: 预测的类别ID, 疾病元组的索引，**从1开始**，顺序为：['normal', 'diabetes', 'glaucoma', 'cataract', 'amd', 'hypertension','myopia', 'other']
            -- Confidence: 预测的置信度，类型为tensor，默认device为cuda
            -- Class Name: 预测的类别名称，类型为str
            -- Inference Time: 推理时间, python默认双精度浮点型，单位为秒
            -- Probabilities : 预测的概率分布，类型为numpy数组，顺序为：['normal', 'diabetes', 'glaucoma', 'cataract', 'amd', 'hypertension','myopia', 'other']
    - Right Eye:
        -- result:
            -- Predicted Class ID: 预测的类别ID, 疾病元组的索引，**从1开始**，顺序为：['normal', 'diabetes', 'glaucoma', 'cataract', 'amd', 'hypertension','myopia', 'other']
            -- Confidence: 预测的置信度，类型为tensor，默认device为cuda
            -- Class Name: 预测的类别名称，类型为str
            -- Inference Time: 推理时间, python默认双精度浮点型，单位为秒
            -- Probabilities : 预测的概率分布，类型为numpy数组，顺序为：['normal', 'diabetes', 'glaucoma', 'cataract', 'amd', 'hypertension','myopia', 'other']
    示例用法：
    left_model_path = best.pt
    right_model_path = best.pt
    left_image_path = '示例左眼.jpg'
    right_image_path = '示例.jpg'
    eye_predictor = EyeDiseasePredictor(left_model_path, right_model_path)
    prediction_result = eye_predictor.diagnose(left_image_path, right_image_path) # 返回上述字典

    # 单独打印预测结果
    print("左眼预测结果：")
    print(prediction_result["Left Eye"]["result"]["Class Name"])
    print("右眼预测结果：")
    print(prediction_result["Right Eye"]["result"]["Class Name"])
    """
    def __init__(self, left_model_path, right_model_path):
        """
        初始化眼底疾病预测器
        :param left_model_path: 左眼模型权重路径
        :param right_model_path: 右眼模型权重路径
        """
        # 加载模型
        self.left_model = YOLO(left_model_path)
        self.right_model = YOLO(right_model_path)

        # 初始化预处理器
        self.clahe_preprocessor = CLAHEPreprocessor(clip_limit=2.0, tile_grid_size=(8, 8))

    def preprocess_images(self, left_image, right_image):
        """
        图像预处理
        :param left_image: 左眼图像输入（可以是文件路径、图像数组或图像文件对象）
        :param right_image: 右眼图像输入（可以是文件路径、图像数组或图像文件对象）
        :return: 预处理后的左右眼图像
        """

        # 辅助函数：读取图像
        def read_image(image_input):
            if isinstance(image_input, str):  # 如果是文件路径
                image = cv2.imread(image_input)
            elif isinstance(image_input, np.ndarray):  # 如果是图像数组
                image = image_input
            else:
                raise ValueError("输入类型不支持，请提供文件路径或图像数组。")
            return image

        # 读取左右眼图像
        left_image = read_image(left_image)
        right_image = read_image(right_image)

        # 检查图像是否加载成功
        if left_image is None:
            raise ValueError(f"左眼图像加载失败，请检查路径或输入：{left_image}")
        if right_image is None:
            raise ValueError(f"右眼图像加载失败，请检查路径或输入：{right_image}")

        # 将图像从 BGR 转换为 RGB
        left_image = cv2.cvtColor(left_image, cv2.COLOR_BGR2RGB)
        right_image = cv2.cvtColor(right_image, cv2.COLOR_BGR2RGB)

        # 应用 CLAHE 对比度增强、亮度平衡和去噪
        pre_left, pre_right = self.clahe_preprocessor(left_image, right_image)

        return pre_left, pre_right

    def predict(self, left_image, right_image):
        """
        对预处理后的左右眼图像进行预测
        :param left_image: 预处理后的左眼图像
        :param right_image: 预处理后的右眼图像
        :return: 左眼和右眼的预测结果 (概率分布和 YOLO 输出)
        """
        # 左眼推理
        start_time = time.time()
        left_output = self.left_model.predict(source=left_image, show=False, save=False)[0]
        end_time = time.time()
        left_inference_time = end_time - start_time

        # 右眼推理
        start_time = time.time()
        right_output = self.right_model.predict(source=right_image, show=False, save=False)[0]
        end_time = time.time()
        right_inference_time = end_time - start_time

        # 获取概率分布
        left_probs = left_output.probs
        right_probs = right_output.probs

        # 获取分类结果
        left_pred_class = left_probs.top1
        left_pred_confidence = left_probs.top1conf
        left_class_name = self.left_model.model.names.get(int(left_pred_class), 'Unknown')

        right_pred_class = right_probs.top1
        right_pred_confidence = right_probs.top1conf
        right_class_name = self.right_model.model.names.get(int(right_pred_class), 'Unknown')

        # 将所有数据封装为字典返回
        result = {
            "Left Eye": {
                "Predicted Class ID": left_pred_class,
                "Confidence": left_pred_confidence,
                "Class Name": left_class_name,
                "Inference Time": left_inference_time,
                "Probabilities": left_probs.data.cpu().numpy()
            },
            "Right Eye": {
                "Predicted Class ID": right_pred_class,
                "Confidence": right_pred_confidence,
                "Class Name": right_class_name,
                "Inference Time": right_inference_time,
                "Probabilities": right_probs.data.cpu().numpy()
            }
        }

        return result

    def diagnose(self, left_image_path, right_image_path):
        """
        对左右眼图像进行预处理和预测，并返回诊断结果
        :param left_image_path: 左眼图像路径
        :param right_image_path: 右眼图像路径
        :return: 诊断结果，包含左右眼的概率分布和YOLO输出
        """
        # 预处理图像
        pre_left, pre_right = self.preprocess_images(left_image_path, right_image_path)

        # 进行预测
        prediction_result = self.predict(pre_left, pre_right)

        # 封装诊断结果
        diagnosis_result = {
            "Left Eye": {
                "result": prediction_result["Left Eye"],
            },
            "Right Eye": {
                "result": prediction_result["Right Eye"],
            },
        }

        return diagnosis_result


if __name__ == "__main__":
    # 示例用法
    left_model_path = '../model/best.pt'
    right_model_path = '../model/best.pt'
    eye_predictor = EyeDiseasePredictor(left_model_path, right_model_path)

    left_image_path = '示例左眼.jpg'
    right_image_path = '示例右眼.jpg'

    # 进行预测
    prediction_result = eye_predictor.diagnose(left_image_path, right_image_path)

    # 单独打印预测结果
    # print(prediction_result)
    print("预测结果：")
    print('左眼：' + str(prediction_result["Left Eye"]['result']['Class Name'])
          + '\n右眼：' + str(prediction_result["Right Eye"]['result']['Class Name']))

